package cn.alex.springaidemo.openAi;

import cn.alex.springaidemo.CommonConstant;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.model.Media;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.FileSystemResource;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.io.File;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 聊天模型
 *
 * @author Alex
 * @date 2024/9/20 00:22
 */
@RestController
@RequestMapping("/chat/model/openAi")
@RequiredArgsConstructor
public class ChatModelController {
    private final ChatModel chatModel;
    private final VectorStore vectorStore;


    /**
     * 一次性全部返回
     */
    @GetMapping("/all")
    public String openAi(@RequestParam(value = "msg", defaultValue = "你是谁？") String msg) {
        ChatResponse call = chatModel.call(
                new Prompt(
                        msg,
                        OpenAiChatOptions.builder()
                                .withModel(OpenAiApi.ChatModel.GPT_3_5_TURBO)
                                .withTemperature(0.8)
                                .build()
                )
        );

        return call.getResult().getOutput().getContent();
    }

    /**
     * 流式返回
     */
    @GetMapping(value = "/stream", produces = "text/event-stream")
    public Flux<ChatResponse> openAiStream(@RequestParam(value = "msg", defaultValue = "你是谁？") String msg) {
        return chatModel.stream(
                new Prompt(
                        msg,
                        OpenAiChatOptions.builder()
                                .withModel(OpenAiApi.ChatModel.GPT_3_5_TURBO)
                                .withTemperature(0.8)
                                .build()
                )
        );
    }


    /**
     * 自定义提示模版
     *
     * @param msg    提问
     * @param aiName ai名字
     */
    @GetMapping("/prompt")
    public String prompt(@RequestParam(value = "msg", defaultValue = "公司法定代表人是谁？") String msg, @RequestParam(value = "aiName", defaultValue = "智联AI") String aiName) {
        String systemText = """
                你是AI小助手。
                你的名字是{aiName}。
                公司简介：统一社会信用代码：914403001922038216，法定代表人：赵明路，注册资本：4,084,113.182万(元)，地址：深圳市龙岗区坂田华为总部办公楼，华为创立于1987年，是全球领先的ICT（信息与通信）基础设施和智能终端提供商。目前华为约有19.5万员工，业务遍及170多个国家和地区，服务全球30多亿人口。 华为致力于把数字世界带入每个人、每个家庭、每个组织，构建万物互联的智能世界：让无处不在的联接，成为人人平等的权利，成为智能世界的前提和基础；为世界提供多样性算力，让云无处不在，让智能无所不及；所有的行业和组织，因强大的数字平台而变得敏捷、高效、生机勃勃；通过AI重新定义体验，让消费者在家居、出行、办公、影音娱乐、运动健康等全场景获得极致的个性化智慧体验。，
                你应该用你的名字，根据公司简介中的内容回复用户的请求，每次回复需要告诉用户你是谁。
                """;

        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemText);
        Message message = systemPromptTemplate.createMessage(Map.of("aiName", aiName));

        Prompt prompt = new Prompt(
                List.of(new UserMessage(msg), message),
                OpenAiChatOptions.builder()
                        .withModel(OpenAiApi.ChatModel.GPT_3_5_TURBO)
                        .withTemperature(0.8)
                        .build()
        );

        List<Generation> results = chatModel.call(prompt).getResults();
        return results.stream().map(item -> item.getOutput().getContent()).collect(Collectors.joining(""));
    }

    /**
     * 多模态
     */
    @GetMapping("multi")
    public String multi(@RequestParam(value = "msg", defaultValue = "解释一下你在这张照片中看到了什么？") String msg) throws Exception {
        UserMessage userMessage = new UserMessage(
                msg,
                new Media(
                        MediaType.IMAGE_PNG,
                        new URL("https://www.google.com/images/branding/googlelogo/1x/googlelogo_color_272x92dp.png")
                )
        );

        Prompt prompt = new Prompt(
                userMessage,
                OpenAiChatOptions.builder()
                        .withModel(OpenAiApi.ChatModel.GPT_4_O)
                        .withTemperature(0.8)
                        .build()
        );
        Generation result = chatModel.call(prompt).getResult();
        return result.getOutput().getContent();
    }

    /**
     * 自定义函数
     */
    @GetMapping("function")
    public String function(@RequestParam(value = "msg", defaultValue = "北京天气怎么样？") String msg) {
        UserMessage userMessage = new UserMessage(msg);

        ChatResponse call = chatModel.call(
                new Prompt(
                        userMessage,
                        OpenAiChatOptions.builder()
                                .withModel(OpenAiApi.ChatModel.GPT_4_O)
                                .withFunction("weatherFunctionService")
                                .build()
                )
        );
        return call.getResult().getOutput().getContent();
    }

    /**
     * 矢量查询
     */
    @GetMapping("vector")
    public List<String> vector(@RequestParam(value = "msg", defaultValue = "北京天气怎么样？") String msg) {
        List<Document> documentList = new ArrayList<>();
        documentList.add(new Document("深圳天气有点热"));
        documentList.add(new Document("北京天气有点冷"));
        documentList.add(new Document("上海天气比较潮湿"));
        documentList.add(new Document("济南天气一会冷一会热"));

        vectorStore.add(documentList);

        SearchRequest searchRequest = SearchRequest.defaults()
                // 关键词
                .withQuery(msg)
                // 查找出相似后返回最大数量
                .withTopK(1)
                // 相似度，0-1，接近1表示完全相同
                .withSimilarityThreshold(0.75);

        List<Document> result = vectorStore.similaritySearch(searchRequest);
        return result.stream().map(Document::getContent).toList();
    }

    /**
     * 创建RAG知识库（添加到向量数据库）
     */
    @GetMapping("createRag")
    public void createRag() {
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(new FileSystemResource(new File(CommonConstant.FILE_STORE_PATH + "Spring介绍.docx")));
        List<Document> read = tikaDocumentReader.read();

        TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();
        List<Document> apply = tokenTextSplitter.apply(read);

        vectorStore.add(apply);
    }
}
